import time
import uuid
from enum import Enum
import logging

from extensions.ext_redis import redis_client

logger = logging.getLogger(__name__)


class ClusterRateLimit:

    UNLIMIT_RATE = -1

    __unlimit_rate = False

    #实际存储时，一个模型有两个key，分别为
    # 模型可用连接数 dify:cluster_rate_limit:{}:{}:count
    # 模型当前活跃连接token set dify:cluster_rate_limit:{}:{}:token
    __ACTIVE_CONNECTION_KEY = "dify:cluster_rate_limit:{}:{}"
    __ACTIVE_CONNECTION_PATTERN = "dify:cluster_rate_limit:*"


    def __connect_record(self, record: str) -> bool:
        return redis_client.sadd(self.__token_key, record) == 1

    def __connect_delete(self, record: str) -> bool:
        return redis_client.srem(self.__token_key, record) == 1

    # def __connect_check(self, record: str) -> bool:
    #     return redis_client.sismember(self.__token_key, record) == 1

    def clean_all_pool(self):
        cursor = 0
        while True:
            cursor, keys = redis_client.scan(cursor=cursor, match=self.__ACTIVE_CONNECTION_PATTERN)
            for key in keys:
                redis_client.delete(key)
            if cursor == 0:
                break
        logger.info("connection pool cleaned")


    class RateLimitType(Enum):
        MODEL = "model"
        APPLICATION = "application"
        @classmethod
        def value_of(cls, value: str):
            for mode in cls:
                if mode.value == value:
                    return mode
            raise ValueError(f"invalid session status type value {value}")

    def __get_redis_key_prefix(self,
                       type: RateLimitType,
                       distinct_name: str):
        # return "temp"
        return self.__ACTIVE_CONNECTION_KEY.format(type.value, distinct_name)

    def __get_redis_count_key(self,
                       type: RateLimitType,
                       distinct_name: str):
        # return "temp"
        return self.__ACTIVE_CONNECTION_KEY.format(type.value, distinct_name) + ":count"

    def __get_redis_token_key(self,
                       type: RateLimitType,
                       distinct_name: str):
        # return "temp"
        return self.__ACTIVE_CONNECTION_KEY.format(type.value, distinct_name) + ":token"

    def __init__(self,
                 type: RateLimitType,
                 distinct_name: str,
                 max_active_connection: int):
        self.key = self.__get_redis_key_prefix(type, distinct_name)
        self.__count_key = self.__get_redis_count_key(type, distinct_name)
        self.__token_key = self.__get_redis_token_key(type, distinct_name)
        self.__init_connect_pool(max_active_connection)

    def __init_connect_pool(self,
                max_active_connection: int) -> bool:
        if max_active_connection == ClusterRateLimit.UNLIMIT_RATE:
            logger.debug("generate unlimit connection pool")
            self.__unlimit_rate = True
            return True
        result = redis_client.setnx(self.__count_key, max_active_connection)
        return result

    # @func_timeout(3)
    def acquire_connection(self, wait_second: int = 3) -> str:
        if (self.__unlimit_rate
                or not redis_client.exists(self.__count_key)):
            return "unlimit_token"
        token = ""

        for i in range(wait_second):
            current_connection = redis_client.decr(self.__count_key)
            if current_connection < 0:
                redis_client.incr(self.__count_key)
            else:
                token = str(uuid.uuid4())
                self.__connect_record(token)
                return token
            time.sleep(1)
        raise TimeoutError()

    def release_connection(self, token: str) -> bool:
        if self.__unlimit_rate:
            return True
        if self.__connect_delete(token):
            redis_client.incr(self.__count_key)
            return True
        else:
            logger.info(f"unknown token {token}")
            return False


class MultipleRateLimit:

    rate_limit_list: list[ClusterRateLimit] = []

    token_limiter_map: dict[str: ClusterRateLimit] = {}

    def __init__(self, *rate_limit: ClusterRateLimit):
        self.rate_limit_list = list(rate_limit)

    def acquire_connection(self) -> bool:
        for i in range(len(self.rate_limit_list)):
            rate_limit = self.rate_limit_list[i]
            try:
                token = rate_limit.acquire_connection()
                logger.debug(f"success to acquire token from {rate_limit.key}: {token}")
                self.token_limiter_map[token] = rate_limit
            except TimeoutError as e:
                logger.info(f"cannot acquire token from {rate_limit.key}, timeout")
                self.release_connection()
                raise TimeoutError
            except Exception as e:
                logger.error(f"failed to acquire token from {rate_limit.key}: {e}")
                self.release_connection()
                raise e
        return True

    def release_connection(self) -> bool:
        result = True
        for token in self.token_limiter_map.keys():
            limiter = self.token_limiter_map[token]
            result = result and limiter.release_connection(token)
        self.token_limiter_map.clear()
        return result



